771C - Bear and Tree Jumps - CodeForces Solution

#include <iostream>
#include <cstring>
#include <algorithm>
#include <cmath>
#include <map>
#include <set>
#include <vector>
#include <queue> 
#include <unordered_set>
#include <unordered_map>
#include <bitset>
#include <stack>
#include <ctime>
#define SEQ 19

using namespace std;

typedef pair<int, int> pii;
typedef pair<double, double> pdd;
typedef pair<int, pii> piii;
typedef pair<int, long long > pil;
typedef long long ll;
typedef unsigned long long ull;
const int N = 400086, MOD = 1e8 + 7, INF = 0x3f3f3f3f, MID = 50000, M = 528;
int dx[4] = {1, 0, -1, 0}, dy[4] = {0, 1, 0, -1};
vector<ll> num;
int n, m, idx, K, top, w[N], root;
int e[N], ne[N], h[N];
bool st[N];
ll res, f[N][8], sum[N], s[N];

int lowbit(int x)
{
    return x & -x;
}

ll gcd(ll a, ll b)
{
    return b ? gcd(b, a % b) : a;
}

inline double rand(double l, double r)
{
    return (double)rand() / RAND_MAX * (r - l) + l;
}

inline ll qmi(ll a, ll b, ll c)
{
    ll res = 1;
    while (b)
    {
        if (b & 1) res = res * a % c;
        a = a * a % c;
        b >>= 1;
    }
    return res;
}

inline ll qmi(ll a, ll b)
{
    ll res = 1;
    while (b)
    {
        if (b & 1) res *= a;
        a *= a;
        b >>= 1;
    }
    return res;
}

inline ll C(ll a, ll b)
{
    if (a < b) return 0;
    ll res = 1;
    for (int i = 1, j = a; i <= b; i++, j--)
    {
        res = res * j % MOD;
        res = res * qmi(i, MOD - 2, MOD) % MOD;
    }
    return res;
}

inline ll get(int len)
{
    if (len <= 0) return 0;
    return (ll)(len + 1) * len / 2;
}

inline int find(ll x)
{
    return lower_bound(num.begin(), num.end(), x) - num.begin();
}

inline void add(int a, int b)
{
    e[idx] = b, ne[idx] = h[a], h[a] = idx++;
}

void init(int r, int fa)
{
    f[r][0] = 1;
    for (int i = h[r]; ~i; i = ne[i])
    {
        int j = e[i];
        if (j == fa) continue;
        init(j, r);
        sum[r] += sum[j] + f[j][0];
        for (int i = 0; i < m - 1; i++) f[r][i + 1] += f[j][i];
        f[r][0] += f[j][m - 1];
    }
}

void dfs(int r, int fa)
{
    for (int i = h[r]; ~i; i = ne[i])
    {
        int j = e[i];
        if (j == fa) continue;
        s[j] = sum[j] + s[r] - sum[j] - f[j][0] + (f[r][0] - f[j][m - 1]);
        int t[7];
        for (int i = 0; i < m; i++) t[i] = f[j][i];
        f[j][0] = t[0] + f[r][m - 1] - t[(m - 2 + m) % m];
        for (int i = 1; i < m; i++) f[j][i] = t[i] + f[r][i - 1] - t[(i - 2 + m) % m];
        dfs(j, r);
    }
    res += s[r];
}

int main()
{
    memset(h, -1, sizeof h);
    cin >> n >> m;
    for (int i = 1; i < n; i++)
    {
        int a, b;
        scanf("%d%d", &a, &b);
        add(a, b), add(b, a);
    }
    init(1, -1);
    s[1] = sum[1];
    dfs(1, -1);
    printf("%lld\n", res >> 1);
    return 0;
}